import torch
from torchvision import datasets, transforms, utils
import random
import matplotlib.pyplot as plt
import numpy as np
from model import Cifar10ResNet
from data_loader import load_cifar10_data

# 设置matplotlib支持中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]

# 类别名称
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse','ship', 'truck')


def visualize_results(model, test_loader, num_images=5):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    model.eval()

    # 创建一个子图网格
    fig, axes = plt.subplots(1, num_images, figsize=(15, 3))
    
    # 随机选择num_images个样本
    indices = list(range(len(test_loader.dataset)))
    random_indices = random.sample(indices, num_images)

    for i, index in enumerate(random_indices):
        image, label = test_loader.dataset[index]
        
        # 将图像和标签准备好用于模型输入
        image_tensor = image.unsqueeze(0).to(device)
        
        with torch.no_grad():
            output = model(image_tensor)
            _, predicted = torch.max(output.data, 1)
        
        # 将图像转换为numpy数组以便显示
        img_np = image.numpy().transpose((1, 2, 0))
        # 反归一化处理
        img_np = img_np * 0.5 + 0.5
        
        # 在子图上显示图像
        axes[i].imshow(img_np)
        axes[i].axis('off')
        
        # 设置标题，显示实际类别和预测类别
        true_class = classes[label]
        pred_class = classes[predicted.item()]
        
        if true_class == pred_class:
            axes[i].set_title(f"实际: {true_class}\n预测: {pred_class}", color='green')
        else:
            axes[i].set_title(f"实际: {true_class}\n预测: {pred_class}", color='red')
    
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    _, test_loader = load_cifar10_data()
    model = Cifar10ResNet()
    
    # 加载训练好的模型
    try:
        model.load_state_dict(torch.load('image_classifier.pth'))
        print("成功加载模型")
    except FileNotFoundError:
        print("未找到模型文件，请确保已经训练并保存了模型")
        exit()
    
    visualize_results(model, test_loader)
